import fitz  # PyMuPDF


class PdfExtractor:
    def __init__(self, path_pdf):
        self.path_pdf = path_pdf
        self.matched_paragraphs = {  # 保存段落
            "3集成电路设计实验结果": {
                "word": [],
            },
            "4集成电路制造实验结果": {
                "word": [],
            },
            "5集成电路封装与测试实验结果": {
                "word": [],
            },
            "6实验结果分析": {
                "word": [],
            },
            "7实验结论": {
                "word": [],
            }
        }
        self.match_images = []  # 保存图片
        self.doc = None
        self.current_title = None

    def load_document(self):
        """加载文档"""
        self.doc = fitz.open(self.path_pdf)

    def extract_paragraphs(self):
        """提取匹配的段落"""
        if self.doc is None:
            self.load_document()

        for page in self.doc:
            text = page.get_text("text")
            text = text.replace(" ", "")
            lines = text.split('\n')
            for line in lines:
                # 检查是否是我们关心的标题
                if line in self.matched_paragraphs:
                    self.current_title = line
                elif self.current_title is not None:
                    # 如果当前有正在处理的标题，将内容添加到相应列表
                    self.matched_paragraphs[self.current_title]["word"].append(line)

    def extract_images(self):
        """从文档中提取图片数据"""
        if self.doc is None:
            self.load_document()

        for page_number in range(len(self.doc)):
            page = self.doc.load_page(page_number)
            images = page.get_images(full=True)
            for img_index, img in enumerate(images):
                xref = img[0]
                base_image = self.doc.extract_image(xref)
                image_bytes = base_image["image"]
                self.match_images.append(image_bytes)

    def get_matched_paragraphs(self):
        """返回匹配的段落字典"""
        return self.matched_paragraphs

    def get_matched_images(self):
        """返回图片列表"""
        return self.match_images


# 使用示例
def test():
    # pdf_extractor = PdfExtractor(
    #     'C://work/【实验室相关】/【2024-10-22】AI判分/10.91.128.2-Tomcat80/202110/mj_3703363_浠ｅ畤_20211027171257.pdf')
    pdf_extractor = PdfExtractor(
        'C://work/【实验室相关】/【2024-10-22】AI判分/10.91.128.2-Tomcat80/202401/mj_5046417_鍟嗗崥闂籣20240103203932.pdf')
    # pdf_extractor = PdfExtractor(
    #     'C://work/【实验室相关】/【2024-10-22】AI判分/10.91.128.2-Tomcat80/202110/mj_3880366_寮犲瓙鐎歘20211015144539.pdf')
    pdf_extractor.extract_paragraphs()
    pdf_extractor.extract_images()
    paragraphs = pdf_extractor.get_matched_paragraphs()
    images = pdf_extractor.get_matched_images()
    print("xxxxxxxxxxxx")


if __name__=="__main__":
    test()